#注意这个代码没有把模型放置到gpu上面有些框架错误
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
from torch.utils.data import DataLoader
import torch
import json

# 1. 加载预训练模型和分词器，为什么选择这个模型可以直接看看内
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model_name = 'bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name,local_files_only=True)
model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2,local_files_only=True)
model.to(device)  

# 2. 创建数据集,也可以使用Dataset库来构建
file=open("data.jsonl","r",encoding="utf-8")
data=file.readlines()

class Dataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data
   
    def __getitem__(self, idx):
        mid=json.loads(self.data[idx])
        texttensor = tokenizer.encode_plus(
            mid["text"],                                      
            add_special_tokens=True,  # 添加特殊标记
            max_length=32,  # 设置最大序列长度
            padding='max_length',  # 进行填充以确保序列达到最大长度
            truncation=True,  # 截断超出最大长度的序列
            return_attention_mask=True,  # 返回注意力掩码
            return_tensors='pt'
        )
        label=mid["label"]
        #   stack expects each tensor to be equal size, but got [1, 9] at entry 0 and [1, 5] at entry 1这个报错应该是dataset的报错
        #   因为torch中构建dataset的对象，tensor大小必须一致，才方便dataset做dataloader等操作。上述错误就是getitem返回的张量大小不一致的报错。
        #   所以我们用tokenizer.encode_plus把短的补长，长的补短。保证张量大小的相同。
        
        return {
            'input_ids': texttensor['input_ids'].flatten(),
            'attention_mask': texttensor['attention_mask'].flatten(),   #0表示发生填充的地方,1表示没有填充的地方
            'labels': torch.tensor(label, dtype=torch.long)
        }
        """
        {'input_ids': tensor([ 101, 4510, 1213,  928, 1384, 2523, 2345,  102,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0]), 'attention_mask': tensor([1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0]), 'labels': tensor(1)}
        这是return的数据格式可以进行微调
        """
        # trainer类的数据格式应该如何构建
    def __len__(self):
        return len(self.data)


dataset=Dataset(data)
print(dataset[0])   #检查数据库的输出

# 3. 设置训练参数
training_args = TrainingArguments(
    output_dir='./SFT_model',          # 输出目录
    num_train_epochs=30,              # 训练轮数
    per_device_train_batch_size=1,  # 每个设备的训练批次大小
    warmup_steps=500,                # 预热步数
    weight_decay=0.01,               # 权重衰减
    logging_dir='./logs',            # 日志目录
    logging_steps=10,
)

# 5. 创建训练器
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,  # dataset的数据如何构建
)

# 6. 训练模型
trainer.train()
# 7. 保存模型
model.save_pretrained('./SFT_model')
tokenizer.save_pretrained('./SFT_model')